// makebases.cc
// Program for creating factor bases for the General Number
// Field Sieve (GNFS).
// It takes a polynomial and creates the necessary bases!
// It is based on the NTL library, which is based on GnuMP

#include <iostream>
#include <fstream>
#include <cstdlib>
#include <ctime>
#include <string>
#include <NTL/ZZ.h>
#include <NTL/RR.h>
#include <NTL/ZZX.h>
#include <NTL/ZZXFactoring.h>
#include <NTL/ZZ_pXFactoring.h>
#include "nfsinput.h"

int FindRootsModP(NTL::ZZX f, NTL::ZZ p, NTL::ZZ* res) {
	int d = NTL::deg(f);
	int num=0;
	NTL::ZZX fmonic = f;
	NTL::ZZ a_d = NTL::coeff(f,d);
	NTL::ZZ tmp =NTL::to_ZZ(1);
	fmonic.rep[d]=NTL::to_ZZ(1);
	
	for(int i=d-1;i>0;i--) {
		tmp *= a_d;
		fmonic.rep[i-1] *= tmp;
	}
	
	NTL::ZZ_p::init(p);
	NTL::ZZ_p a_d_modp = NTL::to_ZZ_p(a_d);
	NTL::ZZ_p factor = NTL::power(a_d_modp,p-NTL::to_ZZ(2));
	NTL::ZZ_p root;
	
	NTL::vec_pair_ZZ_pX_long factors;
	NTL::ZZ_pX FmodP = NTL::to_ZZ_pX(fmonic);
	NTL::berlekamp(factors, FmodP);
	for (int i=0;i<factors.length();i++) {
		if(NTL::deg(factors[i].a)== 1) {
			root = -factors[i].a.rep[0];
			root*=factor;
			res[num]=NTL::rep(root);
			num++;
		}
	}
	return num;
}

NTL::ZZ MakeRationalBase(char* filename, char* primes, NTL::ZZ bound, NTL::ZZ root) {
	// Generates the rational factor base
	// which consists of the first t primes
	NTL::ZZ elements = NTL::to_ZZ(0);
	NTL::ZZ current;
	std::ofstream outfile(filename);
	std::ifstream infile(primes);

	infile >> current;
	
	outfile << current;
	outfile << std::endl << root%current;
	
	while(current < bound) {
		infile >> current;
		outfile << std::endl << current ;
		outfile << std::endl << root%current;
		elements++;
	}
	infile.close();
	outfile.close();
	return elements;
}


NTL::ZZ MakeAlgebraicBase(NTL::ZZX f, char* algebraicBase, char* primes,NTL::ZZ elements) {
	// Generates the algebraic factor base from
	// a polynomial and a rational factor base.

	NTL::ZZ NumElements=NTL::to_ZZ(0);
	int degree=NTL::deg(f);
	int numRoots;

	NTL::ZZ i,p;
	NTL::ZZ zero = NTL::to_ZZ(0);

	NTL::ZZ* roots = new NTL::ZZ[degree];
	
	std::ifstream infile(primes);
	std::ofstream outfile(algebraicBase);
	infile >> p;
	while(p < elements){
		numRoots = FindRootsModP(f,p,roots);
		if(NTL::coeff(f,0)%p==0) {
			outfile << p << std::endl << 0 << std::endl;
			NumElements++;
		}
		for(int i=0; i<numRoots;i++){
			if(roots[i]!=NTL::to_ZZ(0)) {
				outfile << p << std::endl 
					<< roots[i] << std::endl;
				NumElements++;
			}
		}
		infile >> p;
	}

	delete[] roots;
	infile.close();
	outfile.close();
	return p;
}

void MakeQuadraticBase(NTL::ZZX f, char* quadraticBase, char* primes, long elements,NTL::ZZ start) {
	// Generates the algebraic factor base from
	// a polynomial and a rational factor base.

	int NumElements=0;
	int degree=NTL::deg(f);
	int numRoots;

	NTL::ZZ i,p;
	NTL::ZZ zero = NTL::to_ZZ(0);

	NTL::ZZ* roots = new NTL::ZZ[degree];
	
	std::ifstream infile(primes);
	std::ofstream outfile(quadraticBase);
	while (p < start) {
		infile >> p;
	}
	
	while(NumElements < elements){
		infile >> p;
		infile >> p;
		infile >> p;
		numRoots = FindRootsModP(f,p,roots);
		if(NTL::coeff(f,0)%p==0) {
			outfile << p << std::endl << 0 << std::endl;
			NumElements++;
		}
		for(int i=0; i<numRoots;i++){
			if(roots[i]!=NTL::to_ZZ(0)) {
				outfile << p << std::endl 
					<< roots[i] << std::endl;
				NumElements++;
			}
		}
	}

	delete[] roots;
	infile.close();
	outfile.close();
}


int main (int count, char** args){
	if(count==1) {
		std::cout << std::endl;
		std::cout << "Usage: ./pgnfs_makebases init-file" << std::endl;
		std::cout << std::endl;
		return 1;
	}
	std::clock_t start,end;
	
	NFS_instance prob = readData(args[1]);
	std::string polyfile=args[1];
	polyfile+=".poly";
	nfsPoly poly = readPoly((char*) polyfile.c_str());

	std::string rationalfile=args[1];
	rationalfile+=".rfb";
	std::string algebraicfile=args[1];
	algebraicfile+=".afb";
	std::string quadraticfile=args[1];
	quadraticfile+=".qcb";
	
	std::string primefile=prob.primefile;
	int qcbnum=NTL::to_long(prob.qcb_num);

	NTL::RR::SetPrecision(NTL::NumBits(prob.n));
	NTL::ZZ primebound = prob.rfb_bound;
	start=std::clock();	
	NTL::ZZ RFBnum=MakeRationalBase((char*) rationalfile.c_str(),(char*) primefile.c_str(),primebound,poly.m);
	end=std::clock();	
	std::cout << "Rational factor base created in file: " << rationalfile << " (time spent: " << (double) (end - start)/CLOCKS_PER_SEC << " secs) with " << RFBnum << " elements" << std::endl;
	start=std::clock();	
	NTL::ZZ largest = MakeAlgebraicBase(poly.f, (char*) algebraicfile.c_str(),(char*) primefile.c_str(),prob.afb_bound);
	end=std::clock();	
	std::cout << "Algebraic factor base has been created in file: " << algebraicfile << " (time spent: " << (double) (end - start)/CLOCKS_PER_SEC << " secs) With " <<  LineCount((char*) algebraicfile.c_str())/2 << " elements" << std::endl;
	start=std::clock();	
	MakeQuadraticBase(poly.f, (char*) quadraticfile.c_str(),(char*) primefile.c_str(),qcbnum,largest);
	end=std::clock();	
	std::cout << "Quadratic character base has been created in file: " << quadraticfile << " (time spent: " << (double) (end - start)/CLOCKS_PER_SEC << " secs) with " << qcbnum << " elements" << std::endl;
}

